/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __NOMINALEMBED_HH__
#define __NOMINALEMBED_HH__

#include <cassert>
#include <cstring>
#include <vector>

#include "bogackishampine.hh"
#include "fastonebigheader.hh"
#include "model.hh"
#include "normalize.hh"
#include "rating.hh"
#include "pianovec.hh"

namespace playerpiano
{
  class NominalEmbed
    {
      private:
        unsigned int n_dims;
        unsigned int n_workers;
        unsigned int n_labels;
        unsigned int n_items;
        float* alpha;
        float* tau;
        float prioralpha_sq_stddev;
        float priortau_sq_stddev;
        unsigned int failcount;
        unsigned int n_tau;
        float sqrt_n_params;

        static inline float
        square (float x)
          {
            return x * x;
          }

      public:
        static inline float
        num_tau (unsigned int n_labels,
                 unsigned int n_dims)
          {
            return ((2 * n_labels - n_dims - 1) * n_dims) / 2;
          }

        template<typename T,
                 typename U>
        class Taufunc
          {
            private:
              T values;
              unsigned int n_dims;
              unsigned int n_labels;
              float tmp;

            public:
              Taufunc (T            _values,
                       unsigned int _n_dims,
                       unsigned int _n_labels)
                : values (_values),
                  n_dims (_n_dims),
                  n_labels (_n_labels),
                  tmp (0)
                {
                }

              U
              operator() (unsigned int i,
                          unsigned int j)
                {
                  unsigned int offset;

                  if (i < n_dims)
                    {
                      offset = (i * (i - 1)) / 2;
                    }
                  else
                    {
                      offset = num_tau (i, n_dims);
                    }

                  return (j < i) ? values[offset + j] : (tmp = 0, tmp);
                }
          };

        NominalEmbed (float*       _alpha,
                      float*       _tau,
                      float        prioralpha_stddev,
                      float        priortau_stddev,
                      unsigned int _n_dims,
                      unsigned int _n_workers,
                      unsigned int _n_labels,
                      unsigned int _n_items)
          : n_dims (_n_dims),
            n_workers (_n_workers),
            n_labels (_n_labels),
            n_items (_n_items),
            alpha (_alpha),
            tau (_tau),
            prioralpha_sq_stddev (square (prioralpha_stddev)),
            priortau_sq_stddev (square (priortau_stddev)),
            failcount (0),
            n_tau (num_tau (n_labels, n_dims)),
            sqrt_n_params (sqrt (n_labels * n_dims + n_tau))
          {
            assert (n_labels > n_dims);
          }

        NominalEmbed (Model* model)
          : n_dims (strtoul (model->lookup<char*> ("--n_dims").first, 0, 0)),
            n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0)),
            n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0)),
            n_items (strtoul (model->lookup<char*> ("--n_items").first, 0, 0)),
            alpha (model->lookup<float*> ("alpha").first),
            tau (model->lookup<float*> ("tau").first),
            prioralpha_sq_stddev (1.0),
            priortau_sq_stddev (1.0),
            failcount (0),
            n_tau (num_tau (n_labels, n_dims)),
            sqrt_n_params (sqrt (n_labels * n_dims + n_tau))
          {
            assert (n_labels > n_dims);
          }

        void
        apply_deferred_prior (unsigned int  w,
                              uint64_t      s,
                              uint64_t      t,
                              float         eta,
                              float         rho)
          {
            if (t > s)
              {
                float powt = fastpow (t, 1.0f - rho);
                float sigma = prioralpha_sq_stddev; 
                float pows = fastpow (s, 1.0f - rho);
                float decay = fastexp (eta * (pows - powt) / 
                                       (n_items * (1.0f - rho) * sigma));

                float* alphaptr = alpha + w * n_dims * n_labels;

                for (unsigned int j = 0; j < n_labels; ++j)
                  {
                    for (unsigned int k = 0; k < n_dims; ++k)
                      {
                        alphaptr[j * n_dims + k] *= decay;
                      }
                  }
              }
          }

        class DFunc
          {
            private:
              const Rating&             rating;
              const std::vector<float>& pz;
              const float*              alpha;
              float                     beta;
              const float*              tau;
              unsigned int              n_labels;
              unsigned int              n_dims;
              unsigned int              n_tau;

            public:
              DFunc (const Rating&             _rating,
                     const std::vector<float>& _pz,
                     const float*              _alpha,
                     float                     _beta,
                     const float*              _tau,
                     unsigned int              _n_labels,
                     unsigned int              _n_dims)
                  : rating (_rating),
                    pz (_pz),
                    alpha (_alpha),
                    beta (_beta),
                    tau (_tau),
                    n_labels (_n_labels),
                    n_dims (_n_dims),
                    n_tau (num_tau (n_labels, n_dims))
                {
                }

              template<typename Expr>
              PianoVec 
              operator() (const float&,
                          const Expr& e) const
                {
                  std::vector<float> p (n_labels);
                  std::vector<float> norm (n_labels);
                  const float* alphaptr = 
                    alpha + rating.worker * n_dims * n_labels;
                  PianoVec rv (n_dims * n_labels + n_tau, 0);
                  Taufunc<const float*, float> taufunc (tau, n_dims, n_labels);
                  Taufunc<const Expr&, float> exprtaufunc (e, n_dims, n_labels);
                  Taufunc<PianoVec&, float&> rvtaufunc (rv, n_dims, n_labels);

                  for (unsigned int z = 0; z < n_labels; ++z)
                    {
                      for (unsigned int r = 0; r < n_labels; ++r)
                        {
                          norm[r] = 0.0;

                          for (unsigned int l = 0; l < n_dims; ++l)
                            {
                              norm[r] += 
                                square (taufunc (z, l)
                                        + exprtaufunc (z, l) 
                                        + alphaptr[z * n_dims + l]
                                        + e[z * n_dims + l + n_tau]
                                        - taufunc (r, l)
                                        - exprtaufunc (r, l)
                                        - alphaptr[r * n_dims + l]
                                        - e[r * n_dims + l + n_tau]);
                            }

                          p[r] = -beta * norm[r];
                        }

                      safe_exponentiate_normalize (p);

                      for (unsigned int k = 0; k < n_labels; ++k)
                        {
                          if (k != z)
                            {
                              float factor = (k == rating.label) ? 1.0 - p[k]
                                                                 : -p[k];

                              factor *= 2.0 * beta * pz[z];

                              for (unsigned int l = 0; l < n_dims; ++l)
                                {
                                  float v = (taufunc (z, l)
                                             + exprtaufunc (z, l) 
                                             + alphaptr[z * n_dims + l]
                                             + e[z * n_dims + l + n_tau]
                                             - taufunc (k, l)
                                             - exprtaufunc (k, l)
                                             - alphaptr[k * n_dims + l]
                                             - e[k * n_dims + l + n_tau]);

                                  rv[z * n_dims + l + n_tau] -= factor * v;
                                  rvtaufunc (z, l) -= factor * v;
                                  rv[k * n_dims + l + n_tau] += factor * v;
                                  rvtaufunc (k, l) += factor * v;
                                }
                            }
                        }
                    }

                  return rv;
                }
          };

        void
        update (uint64_t                  t,
                float                     eta,
                float                     rho,
                float                     beta,
                const std::vector<float>& pz,
                const Rating&             rating)
          {
            DFunc dfunc (rating, pz, alpha, beta, tau, n_labels, n_dims);
            float etat = eta * fastpow (t, -rho);
            PianoVec s (n_dims * n_labels + n_tau, 0);
            
            std::pair<bool, PianoVec> step = 
              bogacki_shampine_2_3 (dfunc,
                                    0.0f,
                                    s,
                                    etat,
                                    1e-3f * sqrt_n_params,
                                    etat / 1000.0,
                                    etat,
                                    1000);

            if (! step.first)
              {
                ++failcount;
              }
            else
              {
                float* alphaptr = alpha + rating.worker * n_dims * n_labels;
                Taufunc<float*, float&> taufunc (tau, n_dims, n_labels);
                Taufunc<const PianoVec&, float> steptaufunc (step.second,
                                                             n_dims,
                                                             n_labels);

                for (unsigned int j = 0; j < n_labels; ++j)
                  {
                    for (unsigned int k = 0; k < n_dims; ++k)
                      {
                        alphaptr[j * n_dims + k] +=
                          step.second[j * n_dims + k + n_tau];
                        taufunc (j, k) += steptaufunc (j, k);
                      }
                  }
              }
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta) const
          {
            const float* alphaptr = alpha + rating.worker * n_dims * n_labels;
            Taufunc<const float*, float> taufunc (tau, n_dims, n_labels);
            float norm[n_labels];

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                norm[i] = 0.0;

                if (i != true_label)
                  {
                    for (unsigned int j = 0; j < n_dims; ++j)
                      {
                        norm[i] += square (taufunc (true_label, j) 
                                           + alphaptr[true_label * n_dims + j]
                                           - taufunc (i, j)
                                           - alphaptr[i * n_dims + j]);
                      }
                  }
              }

            unsigned int maxindex = 0;
            float max = -beta * norm[0];

            for (unsigned int i = 1; i < n_labels; ++i)
              {
                float v = -beta * norm[i];

                if (v > max)
                  {
                    maxindex = i;
                    max = v;
                  }
              }

            float denominator = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float v = -beta * norm[i];
                float d = v - max;
                float expd = fastexp (d);
                denominator += expd;
              }

            float v = -beta * norm[rating.label];

            return v - max - fastlog (denominator);
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta,
                    float&          dlogbeta,
                    float&          ddlogbeta) const
          {
            const float* alphaptr = alpha + rating.worker * n_dims * n_labels;
            Taufunc<const float*, float> taufunc (tau, n_dims, n_labels);
            float norm[n_labels];

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                norm[i] = 0.0;

                if (i != true_label)
                  {
                    for (unsigned int j = 0; j < n_dims; ++j)
                      {
                        norm[i] += square (taufunc (true_label, j) 
                                           + alphaptr[true_label * n_dims + j]
                                           - taufunc (i, j)
                                           - alphaptr[i * n_dims + j]);
                      }
                  }
              }

            unsigned int maxindex = 0;
            float max = -beta * norm[0];

            for (unsigned int i = 1; i < n_labels; ++i)
              {
                float v = -beta * norm[i];

                if (v > max)
                  {
                    maxindex = i;
                    max = v;
                  }
              }

            float denominator = 0;
            float ddenominatordlogbeta = 0;
            float dddenominatorddlogbeta = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float v = -beta * norm[i];
                float d = v - max;
                float expd = fastexp (d);
                denominator += expd;
                ddenominatordlogbeta += expd * d;
                dddenominatorddlogbeta += expd * d * (d + 1.0);
              }

            float v = -beta * norm[rating.label];

            dlogbeta = v - max - ddenominatordlogbeta / denominator;
            ddlogbeta = v 
                        - max
                        + square (ddenominatordlogbeta / denominator)
                        - dddenominatorddlogbeta / denominator;

            return v - max - fastlog (denominator);
          }


        static void
        ensure (Model* model)
          {
            unsigned int n_dims (strtoul (model->lookup<char*> ("--n_dims").first, 0, 0));
            unsigned int n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0));
            unsigned int n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0));
            unsigned int n_tau (num_tau (n_labels, n_dims));

            if (! model->lookup<void*> ("tau").first)
              {
                model->ensure ("tau", n_tau * sizeof (float));

                float* tau = model->lookup<float*> ("tau").first;

                srand48 (strtoul (model->lookup<char*> ("--seed").first, 0, 0));

                for (unsigned int i = 0; i < n_tau; ++i)
                  {
                    if (2 * i <= n_dims * (n_dims - 1))
                      {
                        // there are reflection symmetries which i'm
                        // not eliminating. :(
                        // the best thing would be to constrain the 
                        // optimization so that these components
                        // are non-negative. hopefully this initialization
                        // puts us into the basin of attraction of
                        // that solution.
                        
                        tau[i] = drand48 () + 1.0;
                      }
                    else
                      {
                        tau[i] = drand48 () - 0.5;
                      }
                  }
              }

            model->ensure ("alpha", 
                           n_dims * n_labels * n_workers * sizeof (float));
          }

        unsigned int
        failures () const
          {
            return failcount;
          }
    };
}

#endif // __NOMINALEMBED_HH__
